monai/transforms/post/dictionary.py (2)

516-652: LGTM with minor improvements needed.

The implementation is solid. A few observations:

1. The _update_spatial_metadata method has complex logic for determining spatial shape from different tensor dimensions. Consider simplifying or adding more explicit comments about the logic.

2. Error messages could be extracted to constants for better maintainability (lines 573, 583, 598, 619-621, 633).

3. The type hints could be more specific - Any is used extensively where more concrete types might be known.

Consider extracting error messages:

+_ERR_HEATMAP_KEYS_LEN = "heatmap_keys length must match keys length."
+_ERR_REF_KEYS_LEN = "ref_image_keys length must match keys length when provided."
+_ERR_SHAPE_LEN = "spatial_shape length must match keys length when providing per-key shapes."
+_ERR_NO_SHAPE = "Unable to determine spatial shape for GenerateHeatmapd. Provide spatial_shape or ref_image_keys."
+_ERR_INVALID_POINTS = "landmark arrays must be 2D or 3D with shape (N, D) or (B, N, D)."
+_ERR_REF_NO_SHAPE = "Reference data must define a shape attribute."

 def _prepare_heatmap_keys(self, heatmap_keys: KeysCollection | None) -> tuple[Hashable, ...]:
     if heatmap_keys is None:
         return tuple(f"{key}_heatmap" for key in self.keys)
     keys_tuple = ensure_tuple(heatmap_keys)
     if len(keys_tuple) == 1 and len(self.keys) > 1:
         keys_tuple = keys_tuple * len(self.keys)
     if len(keys_tuple) != len(self.keys):
-        raise ValueError("heatmap_keys length must match keys length.")
+        raise ValueError(_ERR_HEATMAP_KEYS_LEN)
     return keys_tuple

---

636-650: Simplify spatial metadata update logic.

The _update_spatial_metadata method has nested conditionals that make it hard to follow. The logic for distinguishing batched 2D from non-batched 3D is particularly complex.

Consider a clearer approach:

 def _update_spatial_metadata(self, heatmap: MetaTensor, reference: MetaTensor) -> None:
     """Update spatial metadata of heatmap based on its dimensions."""
-    # Update spatial_shape metadata based on heatmap dimensions
-    if heatmap.ndim == 5:  # 3D batched: (B, C, H, W, D)
-        heatmap.meta["spatial_shape"] = tuple(int(v) for v in heatmap.shape[2:])
-    elif heatmap.ndim == 4:  # 2D batched (B, C, H, W) or 3D non-batched (C, H, W, D)
-        # Need to check if this is batched 2D or non-batched 3D
-        if len(heatmap.shape[1:]) == len(reference.meta.get("spatial_shape", [])):
-            # Non-batched 3D
-            heatmap.meta["spatial_shape"] = tuple(int(v) for v in heatmap.shape[1:])
-        else:
-            # Batched 2D
-            heatmap.meta["spatial_shape"] = tuple(int(v) for v in heatmap.shape[2:])
-    else:  # 2D non-batched: (C, H, W)
-        heatmap.meta["spatial_shape"] = tuple(int(v) for v in heatmap.shape[1:])
+    # Determine if batched based on reference's batch dimension
+    ref_is_batched = len(reference.shape) > len(reference.meta.get("spatial_shape", [])) + 1
+    
+    if heatmap.ndim == 5:  # 3D batched
+        spatial_shape = heatmap.shape[2:]
+    elif heatmap.ndim == 4:
+        # Disambiguate: 2D batched vs 3D non-batched
+        spatial_shape = heatmap.shape[2:] if ref_is_batched else heatmap.shape[1:]
+    else:  # ndim == 3, 2D non-batched
+        spatial_shape = heatmap.shape[1:]
+    
+    heatmap.meta["spatial_shape"] = tuple(int(v) for v in spatial_shape)

tests/transforms/test_generate_heatmapd.py (3)

101-101: Remove unused parameters.

Parameters expected_dtype and uses_ref are not used in the test method.

-def test_dict_with_reference_meta(self, _, points, params, expected_shape, expected_dtype, uses_ref):
+def test_dict_with_reference_meta(self, _, points, params, expected_shape, *_unused):

---

151-151: Remove unused parameter.

Parameter expected_dtype is not used.

-def test_dict_batched_with_ref(self, _, points, params, expected_shape, expected_dtype):
+def test_dict_batched_with_ref(self, _, points, params, expected_shape, _expected_dtype):

---

205-229: Document current behavior limitation.

The test acknowledges that MetaTensor points may inherit incorrect affine. This should be tracked as a known issue.

Should I create an issue to track the affine inheritance behavior when using MetaTensor points with reference images?

monai/transforms/post/array.py (3)

753-893: Well-implemented transform with room for minor improvements.

The GenerateHeatmap transform is well-structured. A few suggestions:

1. The backend class attribute should be annotated with ClassVar (line 769).
2. Consider extracting error messages to constants for maintainability.
3. The _evaluate_gaussian method could benefit from a brief docstring.

Fix the class attribute annotation:

+from typing import ClassVar
 
 class GenerateHeatmap(Transform):
-    backend = [TransformBackends.NUMPY, TransformBackends.TORCH]
+    backend: ClassVar[list] = [TransformBackends.NUMPY, TransformBackends.TORCH]

---

863-865: Add boundary check optimization.

The _is_inside method could short-circuit on first failure.

 @staticmethod
 def _is_inside(center: Sequence[float], bounds: tuple[int, ...]) -> bool:
-    return all(0 <= c < size for c, size in zip(center, bounds))
+    for c, size in zip(center, bounds):
+        if not (0 <= c < size):
+            return False
+    return True

---

881-892: Add docstring for clarity.

The _evaluate_gaussian method would benefit from documentation.

 def _evaluate_gaussian(self, coord_shifts: tuple[torch.Tensor, ...], sigma: tuple[float, ...]) -> torch.Tensor:
+    """Evaluate Gaussian at given coordinate shifts with specified sigmas.
+    
+    Args:
+        coord_shifts: Per-dimension coordinate offsets from center.
+        sigma: Per-dimension standard deviations.
+    
+    Returns:
+        Gaussian values at the specified coordinates.
+    """
     device = coord_shifts[0].device

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

Knowledge base: Disabled due to Reviews -> Disable Knowledge Base setting

monai/transforms/post/array.py (5)

754-769: Clarify coordinate order and output shape in docstring.

Please document the expected coordinate order (e.g., ZYX/XY), units (voxels), and output shapes:

• Non-batched: (N, H, W[, D])
• Batched: (B, N, H, W[, D])

 class GenerateHeatmap(Transform):
-    """
-    Generate per-landmark gaussian response maps for 2D or 3D coordinates.
-    ...
-    """
+    """
+    Generate per-landmark Gaussian heatmaps for 2D or 3D coordinates.
+
+    Notes:
+        - Coordinates are interpreted in voxel units and expected in (Y, X) for 2D or (Z, Y, X) for 3D.
+        - Output shape:
+            - Non-batched points (N, D): (N, H, W[, D])
+            - Batched points (B, N, D): (B, N, H, W[, D])
+        - Each channel corresponds to one landmark.
+
+    Args:
+        sigma: ...
+        ...
+
+    Raises:
+        ValueError: when ``sigma`` is non-positive or ``spatial_shape`` cannot be resolved.
+    """

---

796-844: Avoid full-channel scans for normalization; compute peak within the updated window.

This reduces O(HWD) scans per point to O(window) while maintaining behavior (one Gaussian per channel).

-                region = heatmap[b_idx, idx][window_slices]
-                gaussian = self._evaluate_gaussian(coord_shifts, sigma)
-                torch.maximum(region, gaussian, out=region)
-                if self.normalize:
-                    max_val = heatmap[b_idx, idx].max()
-                    if max_val.item() > 0:
-                        heatmap[b_idx, idx] /= max_val
+                region = heatmap[b_idx, idx][window_slices]
+                gaussian = self._evaluate_gaussian(coord_shifts, sigma)
+                updated = torch.maximum(region, gaussian)
+                # write back
+                region.copy_(updated)
+                if self.normalize:
+                    peak = updated.max()
+                    if peak.item() > 0:
+                        heatmap[b_idx, idx] /= peak

---

871-907: Compute Gaussian in float32 for numerical stability, then cast to target dtype.

Half precision can underflow/overflow for the exponent. Do the math in float32 and cast at the end.

-            coord_shifts.append(torch.arange(start, stop, device=device, dtype=self.torch_dtype) - float(c))
+            coord_shifts.append(torch.arange(start, stop, device=device, dtype=torch.float32) - float(c))

-    def _evaluate_gaussian(self, coord_shifts: tuple[torch.Tensor, ...], sigma: tuple[float, ...]) -> torch.Tensor:
+    def _evaluate_gaussian(self, coord_shifts: tuple[torch.Tensor, ...], sigma: tuple[float, ...]) -> torch.Tensor:
         """
         Evaluate Gaussian at given coordinate shifts with specified sigmas.
         ...
         """
-        device = coord_shifts[0].device
-        shape = tuple(len(axis) for axis in coord_shifts)
-        if 0 in shape:
-            return torch.zeros(shape, dtype=self.torch_dtype, device=device)
-        exponent = torch.zeros(shape, dtype=self.torch_dtype, device=device)
+        device = coord_shifts[0].device
+        shape = tuple(len(axis) for axis in coord_shifts)
+        if 0 in shape:
+            return torch.zeros(shape, dtype=self.torch_dtype, device=device)
+        exponent = torch.zeros(shape, dtype=torch.float32, device=device)
         for dim, (shift, sig) in enumerate(zip(coord_shifts, sigma)):
-            scaled = (shift / float(sig)) ** 2
+            shift32 = shift.to(torch.float32)
+            scaled = (shift32 / float(sig)) ** 2
             reshape_shape = [1] * len(coord_shifts)
             reshape_shape[dim] = shift.numel()
             exponent += scaled.reshape(reshape_shape)
-        return torch.exp(-0.5 * exponent)
+        gauss = torch.exp(-0.5 * exponent)
+        return gauss.to(dtype=self.torch_dtype)

---

845-856: Permit integer spatial_shape broadcasting in the call doc and error text.

Consider clarifying in error text that a single int is broadcast across dims (already implemented).

-                raise ValueError("spatial_shape length must match spatial dimension of the landmarks.")
+                raise ValueError(
+                    "spatial_shape length must match the landmarks' spatial dims (or pass a single int to broadcast)."
+                )

---

772-795: Ruff TRY003: shorten exception messages or centralize them.

Minor lint: several raises have relatively long string literals. You can shorten or move to constants if you care about TRY003.

tests/transforms/test_generate_heatmapd.py (3)

56-56: Use tuple unpacking for clarity.

Prefer (1, *shape) over (1,) + shape.

-            (1,) + shape,
+            (1, *shape),

---

168-171: Generalize max over spatial dims without hardcoding.

Flatten spatial dims to compute per-(B,C) maxima.

-        max_vals = heatmap.max(dim=2)[0].max(dim=2)[0].max(dim=2)[0]
-        np.testing.assert_allclose(
-            max_vals.cpu().numpy(), np.ones((expected_shape[0], expected_shape[1])), rtol=1e-5, atol=1e-5
-        )
+        hm2 = heatmap.reshape(heatmap.shape[0], heatmap.shape[1], -1)
+        max_vals = hm2.max(dim=2)[0]
+        np.testing.assert_allclose(max_vals.cpu().numpy(), np.ones((expected_shape[0], expected_shape[1])), rtol=1e-5, atol=1e-5)

---

205-229: Align affine to the reference image when provided.

Test name and comment imply ref affine should win. Recommend adjusting the transform to explicitly set heatmap.affine = reference.affine when a ref MetaTensor is used. See suggested change in dictionary.py.

monai/transforms/post/dictionary.py (4)

532-551: Broaden dtype parameter type to accept torch.dtype.

Matches the array transform signature and current tests.

-        dtype: np.dtype | type = np.float32,
+        dtype: np.dtype | torch.dtype | type = np.float32,

---

552-573: Ensure affine alignment to reference MetaTensor.

Explicitly copy the affine when reference is MetaTensor, in addition to spatial metadata.

             if reference is not None and isinstance(reference, (torch.Tensor, np.ndarray)):
                 # Convert to match reference type and device while preserving heatmap's dtype
                 heatmap, _, _ = convert_to_dst_type(
                     heatmap, reference, dtype=heatmap.dtype, device=getattr(reference, "device", None)
                 )
                 # Copy metadata if reference is MetaTensor
                 if isinstance(reference, MetaTensor) and isinstance(heatmap, MetaTensor):
+                    heatmap.affine = reference.affine
                     self._update_spatial_metadata(heatmap, reference)

---

516-523: Docstring: add output shapes and per-key behavior.

Document the default for heatmap_keys, shape inference precedence (static_shape > ref_image), and output shapes for batched/non-batched.

---

641-656: Simplify spatial_shape derivation.

You can reliably set spatial_shape as the trailing spatial dims regardless of batch presence:

-        ref_spatial_shape = reference.meta.get("spatial_shape", [])
-        ref_is_batched = len(reference.shape) > len(ref_spatial_shape) + 1
-
-        if heatmap.ndim == 5:  # 3D batched: (B, C, H, W, D)
-            spatial_shape = heatmap.shape[2:]
-        elif heatmap.ndim == 4:  # 2D batched (B, C, H, W) or 3D non-batched (C, H, W, D)
-            # Disambiguate: 2D batched vs 3D non-batched
-            spatial_shape = heatmap.shape[2:] if ref_is_batched else heatmap.shape[1:]
-        else:  # 2D non-batched: (C, H, W)
-            spatial_shape = heatmap.shape[1:]
+        # trailing dims after channel are spatial regardless of batch presence
+        spatial_shape = heatmap.shape[-(reference.ndim - 1):]

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

Knowledge base: Disabled due to Reviews -> Disable Knowledge Base setting

tests/transforms/test_generate_heatmapd.py (4)

121-130: Add dtype assertions for static-shape cases.

You pass expected_dtype but only assert equality with numpy’s dtype for static cases; consider also asserting no NaN/Inf and max == 1.0 for parity with other tests.

---

131-136: Verify error message content.

Also assert the exception string contains a helpful hint (e.g., “Provide spatial_shape or ref_image_keys”) to lock behavior.

---

186-205: Add negative/mismatch cases for key lengths.

Please add tests for:

• Mismatched heatmap_keys length → ValueError.
• Mismatched ref_image_keys length → ValueError.
• Per-key spatial_shape length mismatch → ValueError.

These hit the explicit validation branches.

---

206-231: Align test with current behavior: affine comes from reference.

The comment suggests the heatmap “may inherit from points”, but the implementation sets heatmap.affine from the reference. Prefer asserting that explicitly.

Apply this diff:

-        # Note: Currently the heatmap may inherit affine from points MetaTensor
-        # This test documents the current behavior
-        # Ideally, the heatmap should use the reference image's affine
+        # Heatmap should inherit affine from the reference image
+        assert_allclose(heatmap.affine, image.affine, type_test=False)

monai/transforms/post/dictionary.py (3)

516-527: Docstring lacks parameter/return details.

Briefly document keys, heatmap_keys/ref_image_keys broadcasting rules, dtype/device behavior, and raised exceptions to match the rest of this module.

---

559-581: Spatial metadata should use the resolved shape (not reference.ndim).

Relying on reference.ndim-1 to infer spatial dims can be brittle with unusual channel layouts/batch dims. You already have the resolved shape; pass it through.

Apply this diff (two locations):

@@
-                if isinstance(reference, MetaTensor) and isinstance(heatmap, MetaTensor):
-                    heatmap.affine = reference.affine
-                    self._update_spatial_metadata(heatmap, reference)
+                if isinstance(reference, MetaTensor) and isinstance(heatmap, MetaTensor):
+                    heatmap.affine = reference.affine
+                    self._update_spatial_metadata(heatmap, shape)

And update the helper:

@@
-    def _update_spatial_metadata(self, heatmap: MetaTensor, reference: MetaTensor) -> None:
-        """Update spatial metadata of heatmap based on its dimensions."""
-        # trailing dims after channel are spatial regardless of batch presence
-        spatial_shape = heatmap.shape[-(reference.ndim - 1) :]
-        heatmap.meta["spatial_shape"] = tuple(int(v) for v in spatial_shape)
+    def _update_spatial_metadata(self, heatmap: MetaTensor, spatial_shape: tuple[int, ...]) -> None:
+        """Set spatial_shape explicitly from resolved shape."""
+        heatmap.meta["spatial_shape"] = tuple(int(v) for v in spatial_shape)

---

624-636: Error path for invalid points: message clarity.

Consider appending the actual ndim found to the error string to help users debug malformed inputs.

monai/transforms/post/array.py (5)

756-775: Clarify coordinate order and output shape semantics.

Docstring is good; suggest explicitly stating that target_shape is (Z, Y, X) for 3D and (Y, X) for 2D, and output layout is channel-first per-point heatmaps (N, …) or (B, N, …).

---

825-846: Avoid device sync on normalization.

peak.item() forces host sync on GPU. Keep math on-device to reduce stalls.

Apply this diff:

-                if self.normalize:
-                    peak = updated.max()
-                    if peak.item() > 0:
-                        heatmap[b_idx, idx] /= peak
+                if self.normalize:
+                    peak = updated.amax()
+                    denom = torch.where(peak > 0, peak, torch.ones_like(peak))
+                    heatmap[b_idx, idx] = heatmap[b_idx, idx] / denom

---

854-867: Spatial shape resolution UX.

Broadcasting a single int to D dims is handy; consider adding that behavior to the class docstring for discoverability.

---

896-920: Gaussian evaluation is correct; minor micro-opts possible.

Current implementation is clear; if needed, precompute reshape shapes outside the loop to shave overhead in hot paths.

---

787-797: Ruff TRY003: long messages in exceptions.

Linters flag long exception strings; optional to shorten or centralize messages, but not a blocker.

Also applies to: 810-817, 857-865, 873-873

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

Knowledge base: Disabled due to Reviews -> Disable Knowledge Base setting

monai/transforms/post/array.py (1)

754-922: Heatmap generation looks correct; tighten normalization and minor polish.

The algorithm and I/O handling are solid. One tweak: normalize using the channel’s amax (post-write) instead of the window’s amax to be robust and clearer.

Apply this diff:

@@
-                if self.normalize:
-                    peak = updated.amax()
-                    denom = torch.where(peak > 0, peak, torch.ones_like(peak))
-                    heatmap[b_idx, idx] = heatmap[b_idx, idx] / denom
+                if self.normalize:
+                    peak = heatmap[b_idx, idx].amax()
+                    denom = torch.where(peak > 0, peak, torch.ones_like(peak))
+                    heatmap[b_idx, idx].div_(denom)

Optional:

• Factor exception messages into class-level constants (like in GenerateHeatmapd) to satisfy Ruff TRY003. As per coding guidelines and static analysis hints.

monai/transforms/post/dictionary.py (1)

516-679: Fix sigma doc and validate static_shape rank

• Doc: change “sequence matching number of points” → “sequence matching the number of spatial dimensions.”
• In _determine_shape, convert points first, compute spatial_dims, then if static_shape is set raise ValueError when len(static_shape) != spatial_dims.

@@ class GenerateHeatmapd(MapTransform):
-        sigma: standard deviation for the Gaussian kernel. Can be a single value or sequence matching number of points.
+        sigma: standard deviation for the Gaussian kernel. Can be a single value or a sequence matching the number of spatial dimensions.

@@ def _determine_shape(
-        if static_shape is not None:
-            return static_shape
+        points_t = convert_to_tensor(points, dtype=torch.float32, track_meta=False)
+        if points_t.ndim not in (2, 3):
+            raise ValueError(f"{self._ERR_INVALID_POINTS} Got {points_t.ndim}D tensor.")
+        spatial_dims = int(points_t.shape[-1])
+        if static_shape is not None:
+            if len(static_shape) != spatial_dims:
+                raise ValueError(f"Provided static spatial_shape has {len(static_shape)} dims; expected {spatial_dims}.")
+            return static_shape
         if ref_key is not None and ref_key in data:
             return self._shape_from_reference(data[ref_key], spatial_dims)
         raise ValueError(self._ERR_NO_SHAPE)

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

Knowledge base: Disabled due to Reviews -> Disable Knowledge Base setting

monai/transforms/post/array.py (3)

778-779: Annotate mutable class attribute backend with ClassVar to satisfy Ruff (RUF012).

Minor typing fix; consistent with linter guidance.

Apply this diff:

-    backend = [TransformBackends.NUMPY, TransformBackends.TORCH]
+    backend: ClassVar[list[TransformBackends]] = [TransformBackends.NUMPY, TransformBackends.TORCH]

And add the import (outside this hunk):

from typing import ClassVar

---

830-847: Avoid unnecessary CPU sync and use torch finiteness checks.

.tolist() on CUDA tensors syncs to CPU. Use torch checks first; convert after.

Apply this diff:

-                center_vals = center.tolist()
-                if not np.all(np.isfinite(center_vals)):
-                    continue
+                if not torch.isfinite(center).all():
+                    continue
+                center_vals = [float(c) for c in center.detach().cpu().tolist()]

---

753-776: Docstring polish and TRY003 linter noise.

Messages are clear but long raises trigger TRY003. Optional: shorten messages or centralize them as class constants.

monai/transforms/post/dictionary.py (5)

553-554: Annotate backend with ClassVar to satisfy RUF012.

Minor typing improvement.

Apply this diff:

-    backend = GenerateHeatmap.backend
+    backend: ClassVar = GenerateHeatmap.backend

And add (outside this hunk) if not present in this module:

from typing import ClassVar

---

583-605: MetaTensor propagation: ensure MetaTensor is kept after conversion.

convert_to_dst_type(..., dtype=heatmap.dtype, device=reference.device) may return a plain Tensor if dtype/device match but meta wasn’t tracked. Safer to explicitly wrap to MetaTensor when reference is MetaTensor.

Apply this diff:

-            if reference is not None and isinstance(reference, (torch.Tensor, np.ndarray)):
+            if reference is not None and isinstance(reference, (torch.Tensor, np.ndarray)):
                 # Convert to match reference type and device while preserving heatmap's dtype
                 heatmap, _, _ = convert_to_dst_type(
                     heatmap, reference, dtype=heatmap.dtype, device=getattr(reference, "device", None)
                 )
                 # Copy metadata if reference is MetaTensor
-                if isinstance(reference, MetaTensor) and isinstance(heatmap, MetaTensor):
-                    heatmap.affine = reference.affine
-                    self._update_spatial_metadata(heatmap, shape)
+                if isinstance(reference, MetaTensor):
+                    if not isinstance(heatmap, MetaTensor):
+                        heatmap = MetaTensor(heatmap, affine=reference.affine, meta=reference.meta.copy())
+                    else:
+                        heatmap.affine = reference.affine
+                    self._update_spatial_metadata(heatmap, shape)

---

648-664: Shorten/centralize error strings (TRY003).

Optional: reduce f-string length or reuse class-level constants to silence linter.

---

516-551: Docstring clarity: unify axis notation and output shape description.

Use consistent Z/Y/X naming across 2D/3D (matches array variant) to avoid confusion with H/W/D wording.

---

583-605: Shape source precedence: warn when both static shape and ref are provided but disagree.

Currently static shape silently wins. Consider logging a warning if both exist and shapes differ.

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

Knowledge base: Disabled due to Reviews -> Disable Knowledge Base setting

monai/transforms/post/array.py (4)

803-805: Future-proof the dtype check.

Use the dtype’s is_floating_point attribute to avoid hard-coding allowed types.

Apply this diff:

-        # Validate that dtype is floating-point for meaningful Gaussian values
-        if self.torch_dtype not in (torch.float16, torch.float32, torch.float64, torch.bfloat16):
-            raise ValueError(f"dtype must be a floating-point type, got {self.torch_dtype}")
+        # Validate that dtype is floating-point for meaningful Gaussian values
+        if not getattr(self.torch_dtype, "is_floating_point", False):
+            raise ValueError("dtype must be a floating-point type (e.g., torch.float32 / np.float32).")

---

829-851: Avoid CPU sync and NumPy in the inner loop (GPU-friendly check).

Replace tolist()/np.isfinite with torch-native checks and precompute bounds tensor once per call. Reduces host-device sync and improves performance on CUDA.

Apply this diff:

         heatmap = torch.zeros((batch_size, num_points, *target_shape), dtype=self.torch_dtype, device=device)
         image_bounds = tuple(int(s) for s in target_shape)
+        bounds_t = torch.as_tensor(image_bounds, device=device, dtype=points_t.dtype)
         for b_idx in range(batch_size):
             for idx, center in enumerate(points_t[b_idx]):
-                center_vals = center.tolist()
-                if not np.all(np.isfinite(center_vals)):
-                    continue
-                if not self._is_inside(center_vals, image_bounds):
-                    continue
-                window_slices, coord_shifts = self._make_window(center_vals, radius, image_bounds, device)
+                if not torch.isfinite(center).all():
+                    continue
+                if not ((center >= 0).all() and (center < bounds_t).all()):
+                    continue
+                # _make_window expects Python floats; pass only when needed
+                window_slices, coord_shifts = self._make_window(center.tolist(), radius, image_bounds, device)
                 if window_slices is None:
                     continue
                 region = heatmap[b_idx, idx][window_slices]
                 gaussian = self._evaluate_gaussian(coord_shifts, sigma)
                 updated = torch.maximum(region, gaussian)
                 # write back
                 region.copy_(updated)
                 if self.normalize:
                     peak = heatmap[b_idx, idx].amax()
                     denom = torch.where(peak > 0, peak, torch.ones_like(peak))
                     heatmap[b_idx, idx].div_(denom)

---

807-857: Missing method docstring for call.

Add a short docstring describing inputs/outputs and shapes to match neighboring transforms.

Apply this diff:

     def __call__(self, points: NdarrayOrTensor, spatial_shape: Sequence[int] | None = None) -> NdarrayOrTensor:
+        """
+        Args:
+            points: landmark coordinates as ndarray/Tensor with shape (N, D) or (B, N, D),
+                ordered as (Y, X) for 2D or (Z, Y, X) for 3D.
+            spatial_shape: spatial size as a sequence or single int (broadcasted). If None, uses
+                the value provided at construction.
+        Returns:
+            Heatmaps with shape (N, *spatial) or (B, N, *spatial), one channel per landmark.
+        Raises:
+            ValueError: if points shape/dimension or spatial_shape is invalid.
+        """

---

778-778: Optional: annotate backend as ClassVar to satisfy lint.

If you follow Ruff’s RUF012, annotate backend as ClassVar. This file has other instances without it, so adopt consistently if you enable the rule globally.

Apply this diff within the class:

-    backend = [TransformBackends.NUMPY, TransformBackends.TORCH]
+    backend = [TransformBackends.NUMPY, TransformBackends.TORCH]  # consider: backend: ClassVar[list[TransformBackends]] = [...]

And add at the top of the file if standardizing:

from typing import ClassVar

Based on learnings

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

Knowledge base: Disabled due to Reviews -> Disable Knowledge Base setting

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

Cache: Disabled due to data retention organization setting

Knowledge base: Disabled due to Reviews -> Disable Knowledge Base setting

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

Cache: Disabled due to data retention organization setting

Knowledge base: Disabled due to Reviews -> Disable Knowledge Base setting

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

Cache: Disabled due to data retention organization setting

Knowledge base: Disabled due to Reviews -> Disable Knowledge Base setting

monai/transforms/post/dictionary.py (2)

54-103: Consider sorting __all__ to keep Ruff happy.

__all__ is still not alphabetically sorted (RUF022). If you want a clean lint run, either sort this list or add an explicit suppression for this rule here.

---

516-713: GenerateHeatmapd logic looks solid; tighten docs and clarify batched-input story.

The transform implementation (key broadcasting, shape resolution including scalar spatial_shape, reference-based shape/metadata, and dtype/device alignment via convert_to_dst_type) looks consistent and correct overall.

A few small follow‑ups:

• The docstring’s 3D output shape description ((N, H, W, D)) doesn’t match GenerateHeatmap’s (N, Z, Y, X) convention; consider rephrasing to “(N, *spatial_shape) (e.g., (N, Y, X) in 2D or (N, Z, Y, X) in 3D)” to avoid confusion.
• The code now explicitly enforces points_t.ndim == 2 and treats inputs as (N, D) only. If support for batched (B, N, D) landmarks is still a goal (per the PR description), this will need coordinated changes in both GenerateHeatmap and GenerateHeatmapd; otherwise, it’d be good to adjust the PR/docs to state that only unbatched (N, D) inputs are supported.
• Since _prepare_shapes and _determine_shape do accept a scalar spatial_shape (e.g., 64) and broadcast it across spatial dims and keys, you might want to call that out explicitly in the docstring under spatial_shape to match the actual behavior.

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

Cache: Disabled due to data retention organization setting

Knowledge base: Disabled due to Reviews -> Disable Knowledge Base setting